"""
数据预处理
"""
import copy
from pathlib import Path

from utils import EN_DICT


def get_train_val(data_path, save_path):
    """
    构造sentence和tag对应的数据
    :param data_path:原始文件所在路径
    :param save_path: 保存文件路径
    :return:
    """
    if not save_path.parent.exists():
        save_path.parent.mkdir(parents=True)
    # 1. 加载原始数据
    data_src_lines = []
    with open(data_path, 'r', encoding='utf-8') as reader:
        for line in reader:
            line = dict(eval(line.strip()))
            data_src_lines.append(line)

    # 2. 遍历数据获取对应的结果
    with open(save_path, 'w', encoding='utf-8') as writer:
         for data in data_src_lines:
            """
            1. 获取文本以及对应的标注
            2. 原始文本中的不可见字符替换成：UNK
            3. 特殊字符的处理：UNK
            """
            data_text = list(data['originalText'].strip().replace('\r\n','🚗').replace(' ','🚗'))
            # data_tag = copy.deepcopy(data_text)
            data_tag = ['0' for _ in data_text]
            data_entities = data['entities']
            for entity in data_entities:
                # 获取当前实体类别
                en_type = entity['label_type']
                # 获取当前实体范围
                start_pos = entity['start_pos']
                end_pos = entity['end_pos']
                num_pos = end_pos - start_pos
                # 获取当前实体标注：B-XXX M-XXX E-XXX S-XXX
                try:
                    en_label = EN_DICT[en_type]
                except KeyError as e:
                    EN_DICT[en_type] = en_type
                    en_label = EN_DICT.get(en_type,en_type)
                # 替换实体
                if num_pos == 1:
                    data_tag[start_pos] = f"S-{en_label}"
                else:
                    data_tag[start_pos] = f"B-{en_label}"
                    data_tag[start_pos+1:end_pos-1] = [f"M-{en_label}" for _ in range(end_pos-start_pos-2)]
                    data_tag[end_pos-1] = f"E-{en_label}"

            # check
            assert len(data_tag) == len(data_text), "生成的标签必须和原始文本长度大小一致"
            # print(EN_DICT)

            # 1. 第一种结构
            # for text, tag in zip(data_text, data_tag):
            #     writer.writelines(f"{text} {tag}\n")
            # writer.writelines("\n")

            # 2. 第二种结构  方便一行一行读取数据训练 但 不容易check对应text和tag
            writer.writelines(f'{" ".join(data_text)}\n')
            writer.writelines(f'{" ".join(data_tag)}\n')

    print(f"文件：{save_path}构建完成")

if __name__ == '__main__':
    get_train_val(
        data_path=Path(r'./datas/training.txt'),
        save_path=Path(r'./datas/sentence_tag/train.txt')
    )
